import sys
import psutil
import time
from colorama import Fore, Style

def print_network_stats():
    previous_stats = {}
    while True:
        current_stats = {}
        interface =  'flannel.1'
        stats = psutil.net_io_counters(pernic=True)['flannel.1']
        # for interface, stats in psutil.net_io_counters(pernic=True).items():
        if interface not in previous_stats:
            previous_stats[interface] = {'stats': stats, 'timestamp': time.time()}
            continue
        
        bytes_sent_diff = stats.bytes_sent - previous_stats[interface]['stats'].bytes_sent
        bytes_recv_diff = stats.bytes_recv - previous_stats[interface]['stats'].bytes_recv

        # 计算时间间隔
        time_diff = time.time() - previous_stats[interface]['timestamp']
        
        sent_rate = bytes_sent_diff / 1024 / 1024 / time_diff if time_diff > 0 else 0
        recv_rate = bytes_recv_diff / 1024 / 1024 / time_diff if time_diff > 0 else 0

        # 如果速率小于1MB/s，转换为KB/s
        if sent_rate < 1:
            sent_rate *= 1024
        if recv_rate < 1:
            recv_rate *= 1024

        print(Fore.GREEN + f"Interface: {interface}" + Fore.BLUE + f" Sent: {sent_rate:.2f}{('MB/s', 'KB/s')[sent_rate >= 1]}" + Fore.RED + f" Received: {recv_rate:.2f}{('MB/s', 'KB/s')[recv_rate >= 1]}")

        # 更新统计数据
        current_stats[interface] = {'stats': stats, 'timestamp': time.time()}
        previous_stats[interface] = current_stats[interface]

        time.sleep(1)

print_network_stats()